import torch.nn as nn
import torch
import torch.nn.functional as F
import math, copy, re
import warnings
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

warnings.simplefilter("ignore")
from torch.utils.data import Dataset, DataLoader

print(torch.__version__)
from torch import optim

Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(Device)


class Embedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super(Embedding, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)

    def forward(self, x):
        out = self.embed(x)
        return out


class PositionalEmbedding(nn.Module):
    def __init__(self, max_seq_len, embed_model_dim):
        super(PositionalEmbedding, self).__init__()
        self.embed_dim = embed_model_dim

        pe = torch.zeros(max_seq_len, self.embed_dim)
        for pos in range(max_seq_len):
            for i in range(0, self.embed_dim, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i) / self.embed_dim)))
                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1)) / self.embed_dim)))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x * math.sqrt(self.embed_dim)
        seq_len = x.size(1)
        x = x + torch.autograd.Variable(self.pe[:, :seq_len], requires_grad=False)
        return x


class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim=512, n_heads=8):
        super(MultiHeadAttention, self).__init__()

        self.embed_dim = embed_dim
        self.n_heads = n_heads
        self.single_head_dim = int(self.embed_dim / self.n_heads)

        self.query_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)
        self.key_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)
        self.value_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)
        self.out = nn.Linear(self.n_heads * self.single_head_dim, self.embed_dim)

    def forward(self, key, query, value, mask=None):
        batch_size = key.size(0)
        seq_length = key.size(1)
        seq_length_query = query.size(1)

        key = key.view(batch_size, seq_length, self.n_heads, self.single_head_dim)
        query = query.view(batch_size, seq_length_query, self.n_heads, self.single_head_dim)
        value = value.view(batch_size, seq_length, self.n_heads, self.single_head_dim)

        k = self.key_matrix(key)
        q = self.query_matrix(query)
        v = self.value_matrix(value)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        k_adjusted = k.transpose(-1, -2)
        product = torch.matmul(q, k_adjusted)

        if mask is not None:
            product = product.masked_fill(mask == 0, float("-1e20"))

        product = product / math.sqrt(self.single_head_dim)
        scores = F.softmax(product, dim=-1)
        scores = torch.matmul(scores, v)

        concat = scores.transpose(1, 2).contiguous().view(batch_size, seq_length_query, self.single_head_dim * self.n_heads)
        output = self.out(concat)

        return output


class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, expansion_factor=4, n_heads=8):
        super(TransformerBlock, self).__init__()

        self.attention = MultiHeadAttention(embed_dim, n_heads)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)

        self.feed_forward = nn.Sequential(
            nn.Linear(embed_dim, expansion_factor * embed_dim),
            nn.ReLU(),
            nn.Linear(expansion_factor * embed_dim, embed_dim)
        )

        self.dropout1 = nn.Dropout(0.1)  # 降低dropout率
        self.dropout2 = nn.Dropout(0.1)

    def forward(self, key, query, value):
        attention_out = self.attention(key, query, value)
        attention_residual_out = attention_out + query
        norm1_out = self.dropout1(self.norm1(attention_residual_out))

        feed_fwd_out = self.feed_forward(norm1_out)
        feed_fwd_residual_out = feed_fwd_out + norm1_out
        norm2_out = self.dropout2(self.norm2(feed_fwd_residual_out))

        return norm2_out


class TransformerEncoder(nn.Module):
    def __init__(self, seq_len, vocab_size, embed_dim, num_layers=2, expansion_factor=4, n_heads=8):
        super(TransformerEncoder, self).__init__()

        self.embedding_layer = Embedding(vocab_size, embed_dim)
        self.positional_encoder = PositionalEmbedding(seq_len, embed_dim)

        self.layers = nn.ModuleList([TransformerBlock(embed_dim, expansion_factor, n_heads) for i in range(num_layers)])

    def forward(self, x):
        embed_out = self.embedding_layer(x)
        out = self.positional_encoder(embed_out)
        for layer in self.layers:
            out = layer(out, out, out)

        return out


class DecoderBlock(nn.Module):
    def __init__(self, embed_dim, expansion_factor=4, n_heads=8):
        super(DecoderBlock, self).__init__()

        self.attention = MultiHeadAttention(embed_dim, n_heads=8)
        self.norm = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(0.1)  # 降低dropout率
        self.transformer_block = TransformerBlock(embed_dim, expansion_factor, n_heads)

    def forward(self, key, query, x, mask):
        attention = self.attention(query, query, query, mask=mask)
        query = self.dropout(self.norm(attention + query))
        out = self.transformer_block(key, query, x)

        return out


class TransformerDecoder(nn.Module):
    def __init__(self, target_vocab_size, embed_dim, seq_len, num_layers=2, expansion_factor=4, n_heads=8):
        super(TransformerDecoder, self).__init__()
        self.word_embedding = nn.Embedding(target_vocab_size, embed_dim)
        self.position_embedding = PositionalEmbedding(seq_len, embed_dim)

        self.layers = nn.ModuleList(
            [
                DecoderBlock(embed_dim, expansion_factor=4, n_heads=8)
                for _ in range(num_layers)
            ]
        )
        self.fc_out = nn.Linear(embed_dim, target_vocab_size)
        self.dropout = nn.Dropout(0.1)  # 降低dropout率

    def forward(self, x, enc_out, mask):
        x = self.word_embedding(x)
        x = self.position_embedding(x)
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(enc_out, x, enc_out, mask)

        # 移除softmax，在计算损失时统一处理
        out = self.fc_out(x)
        return out


class Transformer(nn.Module):
    def __init__(self, embed_dim, src_vocab_size, target_vocab_size, seq_length, num_layers=2, expansion_factor=4, n_heads=8):
        super(Transformer, self).__init__()

        self.target_vocab_size = target_vocab_size
        self.encoder = TransformerEncoder(seq_length, src_vocab_size, embed_dim, num_layers=num_layers,
                                          expansion_factor=expansion_factor, n_heads=n_heads)
        self.decoder = TransformerDecoder(target_vocab_size, embed_dim, seq_length, num_layers=num_layers,
                                          expansion_factor=expansion_factor, n_heads=n_heads)

    def make_trg_mask(self, trg):
        batch_size, trg_len = trg.shape
        trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
            batch_size, 1, trg_len, trg_len
        )
        return trg_mask.to(Device)

    def decode(self, src, trg):
        trg_mask = self.make_trg_mask(trg)
        enc_out = self.encoder(src)
        out_labels = []
        batch_size, seq_len = src.shape[0], src.shape[1]
        out = trg
        for i in range(seq_len):
            out = self.decoder(out, enc_out, trg_mask)
            out = out[:, -1, :]
            out = out.argmax(-1)
            out_labels.append(out.item())
            out = torch.unsqueeze(out, axis=0)
        return out_labels

    def forward(self, src, trg):
        trg_mask = self.make_trg_mask(trg)
        enc_out = self.encoder(src)
        outputs = self.decoder(trg, enc_out, trg_mask)
        return outputs


class SimpleDataset(Dataset):
    def __init__(self, num):
        prefix = torch.tensor([0])  # 前面插入0
        suffix = torch.tensor([9])  # 后面追加9

        self.data = [torch.randint(1, 9, size=(5,)) for _ in range(num)]
        self.labels = [torch.cat([prefix, _, suffix], dim=0) for _ in copy.deepcopy(self.data)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]


def create_data(batch_size=5, num=100):
    dataset = SimpleDataset(num)
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader


def train():
    src_vocab_size = 10
    target_vocab_size = 10
    num_layers = 6  # 减少层数，更容易训练
    seq_length = 7  # 与数据实际长度匹配（5+2）

    # 移动模型到设备
    model = Transformer(embed_dim=128,  # 减小维度，更容易训练
                        src_vocab_size=src_vocab_size,
                        target_vocab_size=target_vocab_size,
                        seq_length=seq_length,
                        num_layers=num_layers,
                        expansion_factor=4,
                        n_heads=4)  # 减少头数
    model.to(Device)

    intervals = 50  # 增加训练轮次
    criterion = nn.CrossEntropyLoss()  # 使用CrossEntropyLoss更合适

    # 使用Adam优化器，学习率减小
    # opt = optim.Adam(model.parameters(), lr=1e-2, betas=(0.9, 0.98), eps=1e-9)
    opt = optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9)
    # 学习率调度器
    # scheduler = optim.lr_scheduler.ReduceLROnPlateau(opt, 'min', patience=3, factor=0.5)

    datas = create_data(32, 1000)  # 增加数据量

    for interval in range(intervals):
        model.train()
        loss_all = 0
        for data, label in datas:
            # 移动数据到设备
            data = data.to(Device)
            label = label.to(Device)

            input_label = label[:, :-1]
            output_label = label[:, 1:]

            opt.zero_grad()  # 先清零梯度
            output = model(data, input_label)

            # 计算损失
            loss = criterion(output.reshape(-1, output.size(-1)), output_label.reshape(-1))
            loss_all += loss.item()

            loss.backward()
            # 梯度裁剪，防止梯度爆炸
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()

        avg_loss = loss_all / len(datas)
        # scheduler.step(avg_loss)  # 根据损失调整学习率
        print(f'第 {interval + 1} 轮, 平均损失: {avg_loss:.6f}, 学习率: {opt.param_groups[0]["lr"]:.8f}')

    model.eval()
    with torch.no_grad():
        # 测试数据
        test_in_datas = [1, 2, 3, 4, 5]
        test_in_data = torch.Tensor(test_in_datas).long().view(1, -1)

        # label = [1,2,3,4,5]
        label = [0]
        label = torch.Tensor(label).long().view(1, -1)

        while 1:
            test_in_data = test_in_data.to(Device)
            label = label.to(Device)

            mask = torch.triu(torch.ones(len(label), len(label)), diagonal=1) == 0
            # mask = mask.view(1,1,1,1)
            mask_self = mask.to(Device)

            output = model(test_in_data, label)
            log_softmax = F.log_softmax(output, dim=-1)
            # print(log_softmax)
            max_value = log_softmax.argmax(dim=-1)
            # label = max_value[-1]
            label = torch.cat([label, max_value[:, -1].view(1, -1)], dim=1)
            print(label)
            if max_value[:, -1].item() == 9:  # 9为结束符
                break




if __name__ == '__main__':
    train()
